import re
import sys
import argparse

def convert_time(time_str):
    """将时间格式转换为标准SRT格式（HH:MM:SS,mmm）"""
    # 统一替换毫秒分隔符为逗号
    time_str = re.sub(r'[,.]', ',', time_str, count=1)
    
    # 处理缺少小时部分的情况
    parts = time_str.split(':')
    if len(parts) == 2:  # 只有分钟和秒
        return f"00:{parts[0]}:{parts[1]}"
    elif len(parts) == 3:  # 完整的小时、分钟、秒
        return time_str
    return time_str  # 其他情况保持原样

def process_vtt(vtt_content):
    """转换VTT内容为SRT格式"""
    srt_lines = []
    block_count = 0
    lines = vtt_content.splitlines()
    i = 0
    
    # 跳过WEBVTT头部
    while i < len(lines) and lines[i].strip() in ('WEBVTT', ''):
        i += 1
    
    # 处理字幕块
    while i < len(lines):
        line = lines[i].strip()
        
        # 跳过空行和注释
        if not line or line.startswith('NOTE'):
            i += 1
            continue
        
        # 检测时间轴行
        if '-->' in line:
            block_count += 1
            srt_lines.append(str(block_count))
            
            # 分割时间轴
            times = re.split(r'\s*-->\s*', line)
            start_time = convert_time(times[0].strip())
            
            # 处理结束时间（可能包含样式信息）
            end_time = convert_time(times[1].split(' ', 1)[0].strip())
            srt_lines.append(f"{start_time} --> {end_time}")
            
            # 添加字幕文本
            i += 1
            while i < len(lines) and lines[i].strip():
                # 移除VTT特定标签
                text = re.sub(r'</?[cvuib]>', '', lines[i])
                srt_lines.append(text)
                i += 1
            srt_lines.append('')  # 添加空行分隔
        else:
            i += 1
    
    return '\n'.join(srt_lines)

def ensure_srt_format(content):
    """确保SRT内容格式正确"""
    # 修复时间格式（确保使用逗号分隔）
    content = re.sub(
        r'(\d{2}:\d{2}:\d{2})[.,](\d{3})', 
        r'\1,\2', 
        content
    )
    
    # 确保字幕块格式正确
    lines = content.splitlines()
    output = []
    i = 0
    
    while i < len(lines):
        line = lines[i].strip()
        if not line:
            i += 1
            continue
            
        # 序号行
        if re.match(r'^\d+$', line):
            output.append(line)
            i += 1
            continue
            
        # 时间轴行
        if '-->' in line:
            times = re.split(r'\s*-->\s*', line)
            if len(times) == 2:
                start = convert_time(times[0].strip())
                end = convert_time(times[1].strip())
                output.append(f"{start} --> {end}")
            else:
                output.append(line)  # 保持原样
            i += 1
            continue
            
        # 文本行
        output.append(line)
        i += 1
    
    return '\n'.join(output)

def main():
    parser = argparse.ArgumentParser(description='字幕文件格式转换与修复工具')
    parser.add_argument('-i', '--input', required=True, help='输入文件路径')
    parser.add_argument('-o', '--output', help='输出文件路径(可选)')
    
    args = parser.parse_args()
    
    # 设置默认输出文件名
    if not args.output:
        if args.input.lower().endswith('.vtt'):
            args.output = args.input.rsplit('.', 1)[0] + '.srt'
        else:
            args.output = args.input.rsplit('.', 1)[0] + '_fixed.srt'
    
    try:
        with open(args.input, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 根据文件类型自动选择处理方式
        if args.input.lower().endswith('.vtt'):
            print("检测到VTT格式，正在进行转换...")
            result = process_vtt(content)
        else:
            print("检测到SRT格式，正在进行格式修复...")
            result = ensure_srt_format(content)
        
        with open(args.output, 'w', encoding='utf-8') as f:
            f.write(result)
        
        print(f"处理成功！结果已保存至: {args.output}")
    
    except Exception as e:
        print(f"处理出错: {str(e)}")
        sys.exit(1)

if __name__ == '__main__':
    main()